# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2022/9/1
# @Author  : moe@miku39.fun
# @File    : videoSpider2.py
# @Software: PyCharm
import asyncio
import os
import platform
import random
import shutil
import sys

import aiofiles
import aiohttp
from tqdm.asyncio import tqdm


class Spider:
    def __init__(self, sess, system):
        self.sess = sess
        self.system = system

    @classmethod
    def mkDir(cls, config, video_info):
        finalDir = os.path.join(config["outputPath"], video_info["title"])
        tempDir = os.path.join(config["tempPath"], video_info["title"])
        if not os.path.exists(finalDir):
            os.makedirs(finalDir)
        if not os.path.exists(tempDir):
            os.makedirs(tempDir)
        return finalDir, tempDir

    @classmethod
    async def isNeedDownload(cls, file, source_blocks, size, bar):
        await file.seek(source_blocks[0])
        block = await file.read(source_blocks[-1] - source_blocks[0])
        blockSize = len(block.strip(b"\x00"))
        bar.update(blockSize)
        if blockSize == size:
            return False
        return blockSize

    @classmethod
    def checkFile(cls, video_path, audio_path):
        if not os.path.exists(video_path):
            open(file=video_path, mode="wb")
        if not os.path.exists(audio_path):
            open(file=audio_path, mode="wb")

    @classmethod
    def getTerminalSize(cls):
        if platform.system() == "Windows":
            width, height = shutil.get_terminal_size()
        else:
            try:
                width, height = os.get_terminal_size(sys.stdin.fileno())
            except (AttributeError, ValueError, OSError):
                try:
                    width, height = os.get_terminal_size(sys.stdout.fileno())
                except (AttributeError, ValueError, OSError):
                    width, height = (80, 25)
        return width, height

    def getTotalSize(self, url):
        response = self.sess.head(url=url).headers
        return int(response["Content-Length"])

    def getBlockRangeDict(self, urls, block_size):
        url = random.choice(urls)
        totalSize = self.getTotalSize(url)
        if block_size == 0:
            return totalSize, [(0, totalSize - 1)]
        blockRangeList = [
            (i, i + block_size - 1) for i in range(0, totalSize, block_size)
        ]
        if totalSize % block_size != 0:
            blockRangeList[-1] = (totalSize // block_size * block_size, totalSize - 1)
        else:
            blockRangeList = [(0, totalSize - 1)]
        return totalSize, blockRangeList

    def getBlockRangeInfo(self, vide_info, block_size):
        videoBlockDict = self.getBlockRangeDict(vide_info["video"]["videoURLs"], block_size)
        audioBlockDict = self.getBlockRangeDict(vide_info["audio"]["audioURLs"], block_size)
        return {
            "videoSize": videoBlockDict[0],
            "audioSize": audioBlockDict[0],
            "videoBlocks": videoBlockDict[-1],
            "audioBlocks": audioBlockDict[-1],
            "totalSize": videoBlockDict[0] + audioBlockDict[0],
        }

    def wrapSingleVideo(self, video_info, config, dir_path, index):
        videoPath = os.path.join(
            dir_path[-1], f"{video_info['title']}-video-{video_info['video']['id']}.mp4"
        )
        audioPath = os.path.join(
            dir_path[-1], f"{video_info['title']}-audio-{video_info['audio']['id']}.mp3"
        )
        blockList = {
            "title": video_info["title"],
            "videoURLs": video_info["video"]["videoURLs"],
            "audioURLs": video_info["audio"]["audioURLs"],
            "videoPath": videoPath,
            "audioPath": audioPath,
            "FinalPath": os.path.join(dir_path[0], f"{video_info['title']}.mp4"),
        }
        blockList.update(self.getBlockRangeInfo(video_info, config["blockSize"]))
        blockList.update({"position": index})
        return blockList

    async def sliceDownload(
        self, urls, file_path, block_size, stream, semaphore, bar, source_blocks
    ) -> bool:
        size = 0
        async with semaphore:
            async with aiofiles.open(
                file=file_path, mode="rb+"
            ) as file, aiohttp.ClientSession(
                headers=self.sess.headers.items(), cookies=self.sess.cookies.items()
            ) as asyncSess:
                downloadFlag = False
                fileName = file_path.split("\\")[-1]
                if (
                    sourceSize := await Spider.isNeedDownload(file, source_blocks, block_size, bar)
                ) is not False:
                    size += sourceSize
                while not downloadFlag:
                    try:
                        url = random.choice(urls)
                        headers = {
                            "Range": f"bytes={size + source_blocks[0]}-{source_blocks[-1]}"
                        }
                        async with asyncSess.get(
                            url=url,
                            headers=headers,
                            timeout=aiohttp.ClientTimeout(connect=5, sock_read=10),
                        ) as response:
                            if stream:
                                while True:
                                    chunk = await response.content.read(2**16)
                                    if not chunk:
                                        break
                                    await file.seek(size + source_blocks[0])
                                    await file.write(chunk)
                                    length = len(chunk)
                                    size += length
                                    bar.update(length)
                            else:
                                chunk = await response.read()
                                await file.seek(size + source_blocks[0])
                                await file.write(chunk)
                                length = len(chunk)
                                size += length
                                bar.update(length)
                            if (source_blocks[-1] - source_blocks[0] + 1) == size:
                                downloadFlag = True
                    except (
                        aiohttp.ClientPayloadError,
                        aiohttp.ClientConnectorError,
                        aiohttp.ServerDisconnectedError,
                    ):
                        await asyncio.sleep(1)
                        f"{fileName} 下载出错，尝试重新连接 ╮(￣▽￣)╭"
                    except asyncio.TimeoutError:
                        f"{fileName} 下载超时，尝试重新连接 ╮(￣▽￣)╭"
                return True

    async def singleVideoDownload(self, video_source_dict, stream, main_semaphore) -> str:
        async with main_semaphore:
            taskList = []
            semaphore = asyncio.Semaphore(10)
            with tqdm(
                desc=f"{video_source_dict['title']}: ",
                total=video_source_dict["totalSize"],
                unit="b",
                position=video_source_dict["position"],
                colour="GREEN",
                leave=False,
                ascii=True,
                unit_scale=True,
                ncols=self.terminalSize[0] - 10,
            ) as bar:
                Spider.checkFile(video_source_dict["videoPath"], video_source_dict["audioPath"])
                taskList.extend(
                    [
                        asyncio.create_task(
                            self.sliceDownload(video_source_dict["videoURLs"], video_source_dict["videoPath"],
                                               video_source_dict["videoSize"], stream, semaphore, bar, each)
                        )
                        for each in video_source_dict["videoBlocks"]
                    ]
                )
                taskList.extend(
                    [
                        asyncio.create_task(
                            self.sliceDownload(video_source_dict["audioURLs"], video_source_dict["audioPath"],
                                               video_source_dict["audioSize"], stream, semaphore, bar, each)
                        )
                        for each in video_source_dict["audioBlocks"]
                    ]
                )
                await asyncio.wait(taskList)
        return f"{video_source_dict['title']} 已下载完成!"

    async def __download(self, video_info, video_source_list, config, banner):
        videoBlockList = []
        videoTaskList = []
        mainSemaphore = asyncio.Semaphore(config["workerNum"])
        dirPath = Spider.mkDir(config, video_info)
        self.terminalSize = Spider.getTerminalSize()
        os.system("clear" if self.system == "posix" else "cls")
        print(banner)
        for index, each in enumerate(video_source_list, 1):
            dic = self.wrapSingleVideo(each, config, dirPath, index)
            videoBlockList.append(dic)
            videoTaskList.append(
                asyncio.create_task(
                    self.singleVideoDownload(dic, config["stream"], mainSemaphore)
                )
            )
        for each in tqdm.as_completed(
            fs=videoTaskList,
            desc=f"{video_info['title']}: ",
            total=len(videoTaskList),
            unit="video",
            position=0,
            colour="YELLOW",
            leave=True,
            ascii=True,
            ncols=self.terminalSize[0] - 10,
        ):
            await each
        return videoBlockList

    def download(self, video_info, video_source_list, config, banner) -> list:
        if self.system == "nt":
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
        loop = asyncio.new_event_loop()
        videoBlockList = loop.run_until_complete(
            self.__download(video_info, video_source_list, config, banner)
        )
        return videoBlockList
